import cv2
import numpy as np

# 添加PaddleOCR依赖
from paddleocr import PaddleOCR

# 将OCR实例化移到函数外部，避免重复创建
ocr = PaddleOCR(
    use_angle_cls=False,
    lang='en',
    use_gpu=False,
    rec_char_type='en',
    rec_model_dir='en_number_mobile_v2.0',
    det_model_dir='en_ppocr_mobile_v2.0_det',
    cls_model_dir=None
)

def preprocess_image(img):
    # 合并图像处理操作为单行
    return cv2.filter2D(
        cv2.fastNlMeansDenoising(
            cv2.cvtColor(img, cv2.COLOR_BGR2GRAY), 
            h=10
        ),
        -1, 
        np.array([[0, -1, 0], [-1, 5, -1], [0, -1, 0]])
    )

def detect_characters(img):
    # 提前转换颜色空间
    rgb_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    result = ocr.ocr(rgb_img, cls=False)
    
    # 简化None检查
    if not result or not isinstance(result[0], list):
        return []

    valid_boxes = []
    for line in result[0]:
        if not line or not isinstance(line, (list, tuple)):
            continue

        # 优化解包逻辑
        def deep_unpack(obj):
            if isinstance(obj, (list, tuple)):
                if len(obj) == 4 and all(isinstance(p, (list, tuple)) and len(p)==2 for p in obj):
                    return obj
                if len(obj) == 2 and isinstance(obj[0], (list, tuple)) and len(obj[0]) == 4:
                    return obj
                for item in obj:
                    if isinstance(item, (list, tuple)):
                        res = deep_unpack(item)
                        if res: return res
            return None

        raw_data = deep_unpack(line)
        if not raw_data or len(raw_data) != 2 or not isinstance(raw_data[1], tuple):
            continue
            
        box, text_info = raw_data
        if not text_info or not isinstance(text_info[0], str):
            continue
            
        text = str(text_info[0]).strip().upper()
        if not text:
            continue

        try:
            # 优化坐标处理
            box_points = np.array(box, dtype=np.float32)
            x_coords = box_points[:, 0].round().astype(int)
            y_coords = box_points[:, 1].round().astype(int)
            
            x1 = max(0, min(x_coords))
            y1 = max(0, min(y_coords))
            x2 = min(img.shape[1], max(x_coords))
            y2 = min(img.shape[0], max(y_coords))
            
            w, h = max(0, x2 - x1), max(0, y2 - y1)
            if w > 5 and h > 5:
                valid_boxes.append(((x1, y1, w, h), text))
                
        except Exception as e:
            print(f"坐标解析失败: {e}")
            continue

    return valid_boxes

def main():
    cap = cv2.VideoCapture(0)
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 360)
    
    actual_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    actual_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    
    # 预计算显示尺寸
    display_size = (actual_width, actual_height)
    
    trackers = cv2.legacy.MultiTracker_create()
    tracked_texts = {}
    frame_counter = 0

    while True:
        ret, frame = cap.read()
        if not ret: break
        
        frame = cv2.resize(frame, display_size)
        
        # 优化检测频率逻辑
        if frame_counter % 100 == 0:
            character_data = detect_characters(frame)
            
            trackers = cv2.legacy.MultiTracker_create()
            tracked_texts.clear()
            
            for (rect, text) in character_data:
                try:
                    x, y, w, h = map(int, rect)
                    # 确保文本不会超出图像边界
                    text_y = max(10, y-10)  # 防止文本跑到图像顶部外面
                    
                    # 绘制识别框
                    cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 255), 3)
                    # 绘制文本（增加背景提高可读性）
                    (text_width, text_height), _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.9, 2)
                    cv2.rectangle(frame, (x, text_y - text_height - 5), 
                                (x + text_width, text_y + 5), (0,0,0), -1)
                    cv2.putText(frame, text, (x, text_y), 
                              cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0,255,0), 2)
                    
                    tracker = cv2.legacy.TrackerCSRT_create()
                    trackers.add(tracker, frame, (x, y, w, h))
                    tracked_texts[id(tracker)] = text
                    
                except Exception as e:
                    print(f"绘制错误: {e}")
                    continue
        
        # 优化跟踪器更新
        success, boxes = trackers.update(frame)
        if success:
            for i, box in enumerate(boxes):
                x, y, w, h = map(int, box)
                # 绘制跟踪框
                cv2.rectangle(frame, (x, y), (x + w, y + h), (255,0,0), 2)
                # 获取并绘制对应的文本
                if i < len(tracked_texts):
                    # 直接使用索引获取文本
                    text = list(tracked_texts.values())[i]
                    text_y = max(10, y-10)
                    (text_width, text_height), _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.9, 2)
                    cv2.rectangle(frame, (x, text_y - text_height - 5), 
                                (x + text_width, text_y + 5), (0,0,0), -1)
                    cv2.putText(frame, text, (x, text_y), 
                              cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0,255,0), 2)

        cv2.imshow('Character Tracker', frame)
        frame_counter += 1
        
        if cv2.waitKey(1) == 27: break

    cap.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()
